package com.aliyun.code83.round2.beans;

import java.util.*;

public class Trie {

    TrieNode root;

    /**
     * Initialize your data structure here.
     */
    public Trie() {
        root = new TrieNode();
    }

    /**
     * Inserts a word into the trie.
     */
    public void insertNode(String word) {
        insertNode(word, root, 0);
    }

    /**
     * 将从index处开始的字串插入到root的子节点中，即将index对应的字符插入到root的子节点中
     *
     * @param word
     * @param root
     * @param index
     */
    private void insertNode(String word, TrieNode root, int index) {
        assert index < word.length() && index > -1;
        char cur = word.charAt(index);
        TreeMap<Character, TrieNode> children = root.children;
        if (null == children) {
            children = new TreeMap<>();
            root.children = children;
        }
        if (!children.containsKey(cur)) {
            children.put(cur, new TrieNode(cur));
        }

        if (index == word.length() - 1) {
            children.get(cur).word = word;
            children.get(cur).result = new LinkedList<>();
            return;
        }
        insertNode(word, children.get(cur), index + 1);
    }

    public void insertWord(String word) {
        insertWord(word, root, 0);
    }

    public void insertWord(String word, TrieNode root, int index) {
        if (index >= word.length()) {
            return;
        }
        char cur = word.charAt(index);
        TreeMap<Character, TrieNode> children = root.children;
        if (null == children || null == children.get(cur)) {
            return;
        }
        if (children.get(cur).word != null) {
            children.get(cur).result.add(word);
        }
        insertWord(word, children.get(cur), index + 1);
    }

    public Map<String, List<String>> getAllResult() {
        return getAllResult(root);
    }

    public Map<String, List<String>> getAllResult(TrieNode root) {
        Map<String, List<String>> result = new HashMap<>();
        if (root.word != null) {
            result.put(root.word, root.result);
        }
        if (root.children != null) {
            for (Map.Entry<Character, TrieNode> characterTrieNodeEntry : root.children.entrySet()) {
                Map<String, List<String>> childResult = getAllResult(characterTrieNodeEntry.getValue());
                result.putAll(childResult);
            }
        }
        return result;
    }

    static class TrieNode {
        char c;
        String word;
        TreeMap<Character, TrieNode> children;
        List<String> result;

        public TrieNode() {
        }

        public TrieNode(char c) {
            this.c = c;
        }
    }
}
